import numpy as np
import matplotlib.pyplot as plt

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def _create_mexican_hat_wavelet(t_range=(-1, 1), n_points=200):
    """创建Mexican Hat小波（具有零阶和一阶消失矩）"""
    t_w = np.linspace(t_range[0], t_range[1], n_points)
    psi = (1 - t_w**2) * np.exp(-t_w**2 / 2)
    psi -= np.mean(psi)  # 强制 ∫ψ(t) dt = 0
    return t_w, psi

def _create_haar_wavelet():
    """创建简化的Haar小波（零阶消失矩）"""
    return np.array([1, -1])

def _verify_vanishing_moments(t_w, psi):
    """验证小波的消失矩特性"""
    zero_moment = np.sum(psi) * (t_w[1] - t_w[0])  # 数值积分
    first_moment = np.sum(t_w * psi) * (t_w[1] - t_w[0])
    
    print(f"零阶消失矩验证: ∫ψ dt = {zero_moment:.6f}")
    print(f"一阶消失矩验证: ∫t·ψ dt = {first_moment:.6f}")
    return zero_moment, first_moment

def _align_signals(t, signal, coeff, mode='same'):
    """对齐信号和小波系数的时间轴，处理不同卷积模式"""
    if mode == 'same':
        # 确保长度一致
        min_len = min(len(t), len(signal), len(coeff))
        return t[:min_len], signal[:min_len], coeff[:min_len]
    elif mode == 'valid':
        # 计算有效卷积的起始位置
        start_idx = (len(signal) - len(coeff)) // 2
        end_idx = start_idx + len(coeff)
        t_aligned = t[start_idx:end_idx]
        signal_aligned = signal[start_idx:end_idx]
        # 确保所有数组长度一致
        min_len = min(len(t_aligned), len(signal_aligned), len(coeff))
        return t_aligned[:min_len], signal_aligned[:min_len], coeff[:min_len]
    elif mode == 'full':
        # 扩展时间轴以匹配完整卷积
        dt = t[1] - t[0]
        pad_left = (len(coeff) - len(signal)) // 2
        t_start = t[0] - pad_left * dt
        t_aligned = np.linspace(t_start, t_start + (len(coeff) - 1) * dt, len(coeff))
        signal_padded = np.pad(signal, (pad_left, len(coeff) - len(signal) - pad_left), mode='constant')
        # 确保所有数组长度一致
        min_len = min(len(t_aligned), len(signal_padded), len(coeff))
        return t_aligned[:min_len], signal_padded[:min_len], coeff[:min_len]
    else:
        raise ValueError(f"不支持的卷积模式: {mode}")

def demonstrate_wavelet_filtering(conv_mode='same'):
    """演示一阶消失矩小波对线性趋势的过滤效果
    
    Args:
        conv_mode: 卷积模式 ('same', 'valid', 'full')
    """
    
    # 1. 构造时间轴和信号组件
    t = np.linspace(0, 1, 500)
    linear_trend = 2 * t
    gaussian_peak = np.exp(-((t - 0.5) / 0.05) ** 2)
    signal = linear_trend + gaussian_peak
    
    # 2. 创建小波
    t_w, psi = _create_mexican_hat_wavelet()
    haar = _create_haar_wavelet()
    
    # 3. 小波变换
    mexican_coeff = np.convolve(signal, psi, mode=conv_mode)
    haar_coeff = np.convolve(signal, haar, mode=conv_mode)
    
    # 4. 对齐信号和系数
    t_aligned, signal_aligned, mexican_coeff = _align_signals(t, signal, mexican_coeff, conv_mode)
    _, _, haar_coeff = _align_signals(t, signal, haar_coeff, conv_mode)
    
    # 确保haar_coeff与其他系数长度一致
    min_len = min(len(t_aligned), len(mexican_coeff), len(haar_coeff))
    t_aligned = t_aligned[:min_len]
    signal_aligned = signal_aligned[:min_len]
    mexican_coeff = mexican_coeff[:min_len]
    haar_coeff = haar_coeff[:min_len]
    
    # 5. 验证消失矩特性
    print(f"=== Mexican Hat小波消失矩验证 (模式: {conv_mode}) ===")
    _verify_vanishing_moments(t_w, psi)
    print(f"信号长度: {len(signal)}, 小波系数长度: {len(mexican_coeff)}")
    print(f"对齐后 - t_aligned: {len(t_aligned)}, signal_aligned: {len(signal_aligned)}, mexican_coeff: {len(mexican_coeff)}")
    print(f"haar_coeff: {len(haar_coeff)}")
    
    # 5. 综合可视化
    plt.figure(figsize=(15, 10))
    
    # 子图1：信号组成分析
    plt.subplot(2, 3, 1)
    # 对原始信号组件也进行对齐，使用相同的对齐逻辑
    _, linear_aligned, _ = _align_signals(t, linear_trend, mexican_coeff, conv_mode)
    _, gaussian_aligned, _ = _align_signals(t, gaussian_peak, mexican_coeff, conv_mode)
    
    plt.plot(t_aligned, linear_aligned, label='线性趋势', linestyle='--', alpha=0.7)
    plt.plot(t_aligned, gaussian_aligned, label='高斯峰值', linestyle=':', alpha=0.7)
    plt.plot(t_aligned, signal_aligned, label='合成信号', linewidth=2)
    plt.legend()
    plt.title('信号组成分析')
    plt.grid(True, alpha=0.3)
    
    # 子图2：Mexican Hat小波形状
    plt.subplot(2, 3, 2)
    plt.plot(t_w, psi, 'b-', linewidth=2)
    plt.title('Mexican Hat小波\n(一阶消失矩)')
    plt.xlabel('时间')
    plt.ylabel('幅值')
    plt.grid(True, alpha=0.3)
    plt.axhline(y=0, color='k', linestyle='-', alpha=0.3)
    
    # 子图3：小波响应
    plt.subplot(2, 3, 3)
    plt.plot(t_aligned, mexican_coeff, 'r-', linewidth=2)
    plt.title('Mexican Hat响应\n(趋势被过滤)')
    plt.xlabel('时间')
    plt.ylabel('响应幅值')
    plt.grid(True, alpha=0.3)
    plt.axhline(y=0, color='k', linestyle='-', alpha=0.3)
    
    # 子图4：不同消失矩对比
    plt.subplot(2, 3, 4)
    plt.plot(t_aligned, haar_coeff, label='Haar(零阶消失矩)', alpha=0.7)
    plt.plot(t_aligned, mexican_coeff, label='Mexican Hat(一阶消失矩)', linewidth=2)
    plt.legend()
    plt.title('不同消失矩对比')
    plt.xlabel('时间')
    plt.ylabel('响应幅值')
    plt.grid(True, alpha=0.3)
    
    # 子图5：不同斜率的趋势过滤效果
    plt.subplot(2, 3, 5)
    slopes = [0, 1, 2, 5]
    for slope in slopes:
        test_signal = slope * t + gaussian_peak
        test_coeff = np.convolve(test_signal, psi, mode=conv_mode)
        t_test, _, test_coeff_aligned = _align_signals(t, test_signal, test_coeff, conv_mode)
        # 确保所有测试都使用相同的时间轴长度
        if len(t_test) == len(test_coeff_aligned):
            plt.plot(t_test, test_coeff_aligned, label=f'斜率={slope}', alpha=0.8)
        else:
            print(f"警告：斜率{slope}的维度不匹配 - t_test: {len(t_test)}, coeff: {len(test_coeff_aligned)}")
            # 使用统一的对齐时间轴
            plt.plot(t_aligned[:len(test_coeff_aligned)], test_coeff_aligned, label=f'斜率={slope}', alpha=0.8)
    plt.legend()
    plt.title('不同线性趋势的过滤效果')
    plt.xlabel('时间')
    plt.ylabel('响应幅值')
    plt.grid(True, alpha=0.3)
    
    # 子图6：原始信号vs过滤后信号
    plt.subplot(2, 3, 6)
    plt.plot(t_aligned, signal_aligned, label='原始信号(含趋势)', alpha=0.7)
    plt.plot(t_aligned, mexican_coeff, label='过滤后信号', linewidth=2)
    plt.legend()
    plt.title('过滤效果总览')
    plt.xlabel('时间')
    plt.ylabel('幅值')
    plt.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()
    
    # 6. 教学要点总结
    print("\n=== 教学要点总结 ===")
    print("• 核心概念：一阶消失矩小波能'看穿'线性趋势，只响应非线性变化")
    print("• 实际应用：信号去趋势、边缘检测、特征提取")
    print("• 数学直觉：小波像'变化探测器'，对平直部分'视而不见'")
    print("• 关键观察：Mexican Hat对高斯峰值有强响应，对线性趋势响应接近零")

if __name__ == "__main__":
    # 演示不同卷积模式
    print("\n=== 演示不同卷积模式的效果 ===")
    
    # 默认使用 'same' 模式，也可以尝试 'valid' 或 'full'
    conv_mode = 'valid'  # 测试 'valid' 模式
    
    print(f"当前使用卷积模式: {conv_mode}")
    print("• 'same': 输出长度与输入相同（默认）")
    print("• 'valid': 输出长度较短，无边界效应")
    print("• 'full': 输出长度较长，包含所有卷积结果")
    
    demonstrate_wavelet_filtering(conv_mode)